import os
import random
import json


# 递归遍历文件夹并筛选出特定的文件类型
def deep_traverse_and_extract(file_path):
    res_list_py = []
    res_list_java = []
    res_list_js = []

    # 需要筛选的文件扩展名
    python_ext = ['.py']
    java_ext = ['.java']
    js_ext = ['.js', '.vue', '.jsx', '.ts', '.tsx']
    valid_extensions = ['.java', '.py', '.js', '.vue', '.jsx', '.ts', '.tsx']

    for root, dirs, files in os.walk(file_path):
        for file in files:
            # 获取文件扩展名
            _, ext = os.path.splitext(file)

            # 如果是符合条件的文件类型
            if ext in valid_extensions:
                full_path = os.path.join(root, file)
                try:
                    with open(full_path, 'r', encoding='utf-8', errors='ignore') as file_git:
                        content = file_git.read()

                        # 获取文件内容的长度
                        content_length = len(content)
                        if ext in python_ext:
                            res_list_py.append(content)
                        if ext in java_ext:
                            res_list_java.append(content)
                        if ext in js_ext:
                            res_list_js.append(content)

                        # if content_length > 200:  # 确保文件内容足够长
                        #     rand_int = random.randint(10, 200)
                        #     # 随机选取一个位置并截取
                        #     start_pos = random.randint(0, content_length - 200)
                        #     end_pos = start_pos + rand_int
                        #
                        #     pre = content[:start_pos]
                        #     mid = content[start_pos:end_pos]  # 截取中间100个字符
                        #     suf = content[end_pos:]
                        #
                        #     res_list.append({
                        #         "instruction": f'{pre}<FILL_ME>{suf}',
                        #         "input": "",
                        #         "output": mid,
                        #         "system": "",
                        #         "history": []
                        #     })
                        #
                        #     with open('./output_jsonl/Q_A.json', 'w')as f:
                        #         json.dump(res_list, f, ensure_ascii=False, indent=4)

                except Exception as e:
                    print(e)
    return res_list_py, res_list_java, res_list_js


# 调用函数并传入根目录
root_directory1 = './output_by_keyword'  # 修改为目标文件夹路径
res_py1, res_java1, res_js1 = deep_traverse_and_extract(root_directory1)

with open('output_jsonl/output_by_keyword_py.json', 'w', encoding='utf-8') as f:
    json.dump(res_py1, f, ensure_ascii=False, indent=4)
with open('output_jsonl/output_by_keyword_java.json', 'w', encoding='utf-8') as f:
    json.dump(res_java1, f, ensure_ascii=False, indent=4)
with open('output_jsonl/output_by_keyword_js.json', 'w', encoding='utf-8') as f:
    json.dump(res_js1, f, ensure_ascii=False, indent=4)
